import numpy as np
import pandas as pd
import random
import time
import os
import datetime

def generate_customer_profiles_table(n_customers, random_state=0):
    """
    生成客户配置表
    
    参数:
        n_customers: 客户数量
        random_state: 随机种子，用于结果可重现
        
    返回:
        包含客户配置信息的DataFrame
    """
    
    np.random.seed(random_state)
        
    customer_id_properties=[]
    
    # 从随机分布生成客户属性
    for customer_id in range(n_customers):
        
        # 客户位置坐标 (0-100范围内随机)
        x_customer_id = np.random.uniform(0,100)
        y_customer_id = np.random.uniform(0,100)
        
        # 平均交易金额 (5-100范围内随机)
        mean_amount = np.random.uniform(5,100) # 任意但合理的值
        std_amount = mean_amount/2 # 标准差设为平均值的一半
        
        # 每天平均交易次数 (0-4范围内随机)
        mean_nb_tx_per_day = np.random.uniform(0,4) # 任意但合理的值 
        
        customer_id_properties.append([customer_id,
                                      x_customer_id, y_customer_id,
                                      mean_amount, std_amount,
                                      mean_nb_tx_per_day])
        
    customer_profiles_table = pd.DataFrame(customer_id_properties, columns=['CUSTOMER_ID',
                                                                      'x_customer_id', 'y_customer_id',
                                                                      'mean_amount', 'std_amount',
                                                                      'mean_nb_tx_per_day'])
    
    return customer_profiles_table

def generate_terminal_profiles_table(n_terminals, random_state=0):
    """
    生成终端配置表
    
    参数:
        n_terminals: 终端数量
        random_state: 随机种子，用于结果可重现
        
    返回:
        包含终端配置信息的DataFrame
    """
    
    np.random.seed(random_state)
        
    terminal_id_properties=[] 
    
    # 从随机分布生成终端属性
    for terminal_id in range(n_terminals):
        
        # 终端位置坐标 (0-100范围内随机)
        x_terminal_id = np.random.uniform(0,100)
        y_terminal_id = np.random.uniform(0,100)
        
        terminal_id_properties.append([terminal_id,
                                      x_terminal_id, y_terminal_id])
                                       
    terminal_profiles_table = pd.DataFrame(terminal_id_properties, columns=['TERMINAL_ID',
                                                                      'x_terminal_id', 'y_terminal_id'])
    
    return terminal_profiles_table

def get_list_terminals_within_radius(customer_profile, x_y_terminals, r):
    """
    获取客户周围指定半径内的终端列表
    
    参数:
        customer_profile: 客户配置信息
        x_y_terminals: 所有终端的坐标数组
        r: 搜索半径
        
    返回:
        在指定半径范围内的终端ID列表
    """
    
    # 使用numpy数组加速计算
    
    # 将客户位置(x,y)转换为numpy数组
    x_y_customer = customer_profile[['x_customer_id','y_customer_id']].values.astype(float)
    
    # 计算客户与终端位置坐标之间的平方差
    squared_diff_x_y = np.square(x_y_customer - x_y_terminals)
    
    # 按行求和并计算平方根得到距离
    dist_x_y = np.sqrt(np.sum(squared_diff_x_y, axis=1))
    
    # 获取距离小于r的终端索引
    available_terminals = list(np.where(dist_x_y<r)[0])
    
    # 返回终端ID列表
    return available_terminals
    
def generate_transactions_table(customer_profile, start_date = "2018-04-01", nb_days = 10):
    """
    为特定客户生成交易表
    
    参数:
        customer_profile: 客户配置信息
        start_date: 开始日期
        nb_days: 天数
        
    返回:
        包含客户交易信息的DataFrame
    """
    
    customer_transactions = []
    
    # 设置随机种子以确保结果可重现
    random.seed(int(customer_profile.CUSTOMER_ID))
    np.random.seed(int(customer_profile.CUSTOMER_ID))
    
    # 遍历每一天
    for day in range(nb_days):
        
        # 为当天生成随机交易次数
        nb_tx = np.random.poisson(customer_profile.mean_nb_tx_per_day)
        
        # 如果交易次数大于0，则生成交易记录
        if nb_tx>0:
            
            for tx in range(nb_tx):
                
                # 交易时间：围绕中午时间，标准差20000秒。这个选择旨在模拟大多数交易发生在白天的事实。
                time_tx = int(np.random.normal(86400/2, 20000))
                
                # 如果交易时间在0到86400秒之间，则保留，否则丢弃
                if (time_tx>0) and (time_tx<86400):
                    
                    # 交易金额从正态分布中抽取
                    amount = np.random.normal(customer_profile.mean_amount, customer_profile.std_amount)
                    
                    # 如果金额为负数，则从均匀分布中重新抽取
                    if amount<0:
                        amount = np.random.uniform(0,customer_profile.mean_amount*2)
                    
                    amount=np.round(amount,decimals=2)
                    
                    # 如果可用终端列表不为空
                    if len(customer_profile.available_terminals)>0:
                        
                        # 随机选择一个可用终端
                        terminal_id = random.choice(customer_profile.available_terminals)
                    
                        customer_transactions.append([time_tx+day*86400, day,
                                                      customer_profile.CUSTOMER_ID, 
                                                      terminal_id, amount])
            
    customer_transactions = pd.DataFrame(customer_transactions, columns=['TX_TIME_SECONDS', 'TX_TIME_DAYS', 'CUSTOMER_ID', 'TERMINAL_ID', 'TX_AMOUNT'])
    
    # 如果有交易记录
    if len(customer_transactions)>0:
        # 将时间戳转换为日期时间格式
        customer_transactions['TX_DATETIME'] = pd.to_datetime(customer_transactions["TX_TIME_SECONDS"], unit='s', origin=start_date)
        # 重新排列列的顺序
        customer_transactions=customer_transactions[['TX_DATETIME','CUSTOMER_ID', 'TERMINAL_ID', 'TX_AMOUNT','TX_TIME_SECONDS', 'TX_TIME_DAYS']]
    
    return customer_transactions  
    
    
def generate_dataset(n_customers = 10000, n_terminals = 1000000, nb_days=90, start_date="2018-04-01", r=5):
    """
    生成完整数据集
    
    参数:
        n_customers: 客户数量
        n_terminals: 终端数量
        nb_days: 天数
        start_date: 开始日期
        r: 客户与终端之间的最大距离
        
    返回:
        (客户配置表, 终端配置表, 交易数据表)的元组
    """
    
    start_time=time.time()
    customer_profiles_table = generate_customer_profiles_table(n_customers, random_state = 0)
    print("生成客户配置表耗时: {0:.2}s".format(time.time()-start_time))
    
    start_time=time.time()
    terminal_profiles_table = generate_terminal_profiles_table(n_terminals, random_state = 1)
    print("生成终端配置表耗时: {0:.2}s".format(time.time()-start_time))
    
    start_time=time.time()
    # 获取所有终端的坐标
    x_y_terminals = terminal_profiles_table[['x_terminal_id','y_terminal_id']].values.astype(float)
    # 为每个客户关联半径范围内的终端
    customer_profiles_table['available_terminals'] = customer_profiles_table.apply(lambda x : get_list_terminals_within_radius(x, x_y_terminals=x_y_terminals, r=r), axis=1)
    # 使用Pandarallel的替代方案（被注释）
    #customer_profiles_table['available_terminals'] = customer_profiles_table.parallel_apply(lambda x : get_list_closest_terminals(x, x_y_terminals=x_y_terminals, r=r), axis=1)
    # 计算每个客户可用的终端数量
    customer_profiles_table['nb_terminals']=customer_profiles_table.available_terminals.apply(len)
    print("为客户关联终端耗时: {0:.2}s".format(time.time()-start_time))
    
    start_time=time.time()
    # 为每个客户生成交易记录
    transactions_df=customer_profiles_table.groupby('CUSTOMER_ID').apply(lambda x : generate_transactions_table(x.iloc[0], nb_days=nb_days)).reset_index(drop=True)
    # 使用Pandarallel的替代方案（被注释）
    #transactions_df=customer_profiles_table.groupby('CUSTOMER_ID').parallel_apply(lambda x : generate_transactions_table(x.iloc[0], nb_days=nb_days)).reset_index(drop=True)
    print("生成交易记录耗时: {0:.2}s".format(time.time()-start_time))
    
    # 按时间顺序对交易记录进行排序
    transactions_df=transactions_df.sort_values('TX_DATETIME')
    # 重置索引，从0开始
    transactions_df.reset_index(inplace=True,drop=True)
    transactions_df.reset_index(inplace=True)
    # TRANSACTION_ID即为DataFrame的索引，从0开始
    transactions_df.rename(columns = {'index':'TRANSACTION_ID'}, inplace = True)
    
    return (customer_profiles_table, terminal_profiles_table, transactions_df)
    
def add_frauds(customer_profiles_table, terminal_profiles_table, transactions_df):
    """
    向数据集中添加欺诈交易
    
    参数:
        customer_profiles_table: 客户配置表
        terminal_profiles_table: 终端配置表
        transactions_df: 交易数据表
        
    返回:
        包含欺诈标签的交易数据表
    """
    
    # 默认情况下，所有交易都是正常交易
    transactions_df['TX_FRAUD']=0
    transactions_df['TX_FRAUD_SCENARIO']=0
    
    # 欺诈场景1：大额交易
    # 将金额超过220的交易标记为欺诈
    transactions_df.loc[transactions_df.TX_AMOUNT>220, 'TX_FRAUD']=1
    transactions_df.loc[transactions_df.TX_AMOUNT>220, 'TX_FRAUD_SCENARIO']=1
    nb_frauds_scenario_1=transactions_df.TX_FRAUD.sum()
    print("欺诈场景1中的欺诈交易数量: "+str(nb_frauds_scenario_1))
    
    # 欺诈场景2：终端被盗用
    # 遍历每一天
    for day in range(transactions_df.TX_TIME_DAYS.max()):
        
        # 每天随机选择2个终端作为被盗用终端
        compromised_terminals = terminal_profiles_table.TERMINAL_ID.sample(n=2, random_state=day)
        
        # 找到在当天及之后28天内使用这些被盗用终端的交易
        compromised_transactions=transactions_df[(transactions_df.TX_TIME_DAYS>=day) & 
                                                    (transactions_df.TX_TIME_DAYS<day+28) & 
                                                    (transactions_df.TERMINAL_ID.isin(compromised_terminals))]
                            
        # 将这些交易标记为欺诈
        transactions_df.loc[compromised_transactions.index,'TX_FRAUD']=1
        transactions_df.loc[compromised_transactions.index,'TX_FRAUD_SCENARIO']=2
    
    nb_frauds_scenario_2=transactions_df.TX_FRAUD.sum()-nb_frauds_scenario_1
    print("欺诈场景2中的欺诈交易数量: "+str(nb_frauds_scenario_2))
    
    # 欺诈场景3：客户账户被盗用
    # 遍历每一天
    for day in range(transactions_df.TX_TIME_DAYS.max()):
        
        # 每天随机选择3个客户作为被盗用客户
        compromised_customers = customer_profiles_table.CUSTOMER_ID.sample(n=3, random_state=day).values
        
        # 找到在当天及之后14天内这些被盗用客户的交易
        compromised_transactions=transactions_df[(transactions_df.TX_TIME_DAYS>=day) & 
                                                    (transactions_df.TX_TIME_DAYS<day+14) & 
                                                    (transactions_df.CUSTOMER_ID.isin(compromised_customers))]
        
        nb_compromised_transactions=len(compromised_transactions)
        
        # 随机选择1/3的交易，将其金额放大5倍并标记为欺诈
        random.seed(day)
        index_fauds = random.sample(list(compromised_transactions.index.values),k=int(nb_compromised_transactions/3))
        
        transactions_df.loc[index_fauds,'TX_AMOUNT']=transactions_df.loc[index_fauds,'TX_AMOUNT']*5
        transactions_df.loc[index_fauds,'TX_FRAUD']=1
        transactions_df.loc[index_fauds,'TX_FRAUD_SCENARIO']=3
        
                             
    nb_frauds_scenario_3=transactions_df.TX_FRAUD.sum()-nb_frauds_scenario_2-nb_frauds_scenario_1
    print("欺诈场景3中的欺诈交易数量: "+str(nb_frauds_scenario_3))
    
    return transactions_df

# n_customers = 5000
# n_terminals = 10000
# terminal_profiles_table = generate_terminal_profiles_table(n_terminals, random_state = 0)
# x_y_terminals = terminal_profiles_table[['x_terminal_id','y_terminal_id']].values.astype(float)
# customer_profiles_table = generate_customer_profiles_table(n_customers, random_state = 0)
# customer_profiles_table['available_terminals']=customer_profiles_table.apply(lambda x : get_list_terminals_within_radius(x, x_y_terminals=x_y_terminals, r=50), axis=1)


# transactions_df=customer_profiles_table.groupby('CUSTOMER_ID').apply(lambda x : generate_transactions_table(x.iloc[0], nb_days=5)).reset_index(drop=True)
(customer_profiles_table, terminal_profiles_table, transactions_df)=\
    generate_dataset(n_customers = 5000, 
                     n_terminals = 10000, 
                     nb_days=183, 
                     start_date="2025-04-01", 
                     r=5)
transactions_df = add_frauds(customer_profiles_table, terminal_profiles_table, transactions_df)

DIR_OUTPUT = "./simulated-data-raw/"

if not os.path.exists(DIR_OUTPUT):
    os.makedirs(DIR_OUTPUT)

start_date = datetime.datetime.strptime("2025-04-01", "%Y-%m-%d")

for day in range(transactions_df.TX_TIME_DAYS.max()+1):
    
    transactions_day = transactions_df[transactions_df.TX_TIME_DAYS==day].sort_values('TX_TIME_SECONDS')
    
    date = start_date + datetime.timedelta(days=day)
    filename_output = date.strftime("%Y-%m-%d")+'.pkl'
    
    # Protocol=4 required for Google Colab
    transactions_day.to_pickle(DIR_OUTPUT+filename_output, protocol=4)                 